/*
 * Copyright (c) 2023 Shenzhen Kaihong Digital Industry Development Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "rtmp_amf.h"
#include <cassert>
#include <cstddef>
#include <cstring>
#include <netinet/in.h>
#include <stdexcept>
#include "utils.h"

namespace OHOS {
namespace Sharing {

//-------------------- AmfValue --------------------//
AmfValue::AmfValue(AmfType type) : _type(type)
{
    Init();
}

AmfValue::AmfValue(const char *s) : _type(AMF_STRING)
{
    Init();
    *_value.string = s;
}

AmfValue::AmfValue(const std::string &s) : _type(AMF_STRING)
{
    Init();
    *_value.string = s;
}

AmfValue::AmfValue(double n) : _type(AMF_NUMBER)
{
    Init();
    _value.number = n;
}

AmfValue::AmfValue(int i) : _type(AMF_INTEGER)
{
    Init();
    _value.integer = i;
}

AmfValue::AmfValue(bool b) : _type(AMF_BOOLEAN)
{
    Init();
    _value.boolean = b;
}

AmfValue::AmfValue(const AmfValue &from) : _type(AMF_NULL)
{
    *this = from;
}

AmfValue &AmfValue::operator=(const AmfValue &from)
{
    Destroy();
    _type = from._type;
    Init();
    switch (_type) {
        case AMF_STRING:
            *_value.string = (*from._value.string);
            break;
        case AMF_OBJECT:
        case AMF_ECMA_ARRAY:
            *_value.object = (*from._value.object);
            break;
        case AMF_STRICT_ARRAY:
            *_value.array = (*from._value.array);
            break;
        case AMF_NUMBER:
            _value.number = from._value.number;
            break;
        case AMF_INTEGER:
            _value.integer = from._value.integer;
            break;
        case AMF_BOOLEAN:
            _value.boolean = from._value.boolean;
            break;
        default:
            break;
    }
    return *this;
}

AmfValue::~AmfValue()
{
    Destroy();
}

void AmfValue::Clear()
{
    switch (_type) {
        case AMF_STRING:
            _value.string->clear();
            break;
        case AMF_OBJECT:
        case AMF_ECMA_ARRAY:
            _value.object->clear();
            break;
        default:
            break;
    }
}

AmfType AmfValue::Type() const
{
    return _type;
}

const std::string &AmfValue::AsString() const
{
    if (_type != AMF_STRING) {
        static std::string empty;
        return empty;
    }
    return *_value.string;
}

double AmfValue::AsNumber() const
{
    switch (_type) {
        case AMF_NUMBER:
            return _value.number;
        case AMF_INTEGER:
            return _value.integer;
        case AMF_BOOLEAN:
            return _value.boolean;
        default:
            return 0.0;
    }
}

int AmfValue::AsInteger() const
{
    switch (_type) {
        case AMF_NUMBER:
            return (int)_value.number;
        case AMF_INTEGER:
            return _value.integer;
        case AMF_BOOLEAN:
            return _value.boolean;
        default:
            return 0;
    }
}

bool AmfValue::AsBoolean() const
{
    switch (_type) {
        case AMF_NUMBER:
            return _value.number;
        case AMF_INTEGER:
            return _value.integer;
        case AMF_BOOLEAN:
            return _value.boolean;
        default:
            return false;
    }
}

std::string AmfValue::ToString() const
{
    switch (_type) {
        case AMF_NUMBER:
            return std::to_string(_value.number);
        case AMF_INTEGER:
            return std::to_string(_value.integer);
        case AMF_BOOLEAN:
            return _value.boolean ? "true" : "false";
        case AMF_STRING:
            return *(_value.string);
        case AMF_OBJECT:
            return "object";
        case AMF_NULL:
            return "null";
        case AMF_UNDEFINED:
            return "undefined";
        case AMF_ECMA_ARRAY:
            return "ecma_array";
        case AMF_STRICT_ARRAY:
            return "strict_array";
        default:
            return {};
    }
}

const AmfValue &AmfValue::operator[](const char *str) const
{
    if (_type != AMF_OBJECT && _type != AMF_ECMA_ARRAY) {
        static AmfValue v;
        return v;
    }
    auto i = _value.object->find(str);
    if (i == _value.object->end()) {
        static AmfValue val(AMF_NULL);
        return val;
    }
    return i->second;
}

void AmfValue::ObjectForEach(const std::function<void(const std::string &, const AmfValue &)> &fun) const
{
    if (_type != AMF_OBJECT && _type != AMF_ECMA_ARRAY) {
        return;
    }
    for (auto &pr : *(_value.object)) {
        fun(pr.first, pr.second);
    }
}

AmfValue::operator bool() const
{
    return _type != AMF_NULL;
}

void AmfValue::Set(const std::string &s, const AmfValue &val)
{
    if (_type != AMF_OBJECT && _type != AMF_ECMA_ARRAY) {
        return;
    }
    _value.object->emplace(s, val);
}
void AmfValue::Add(const AmfValue &val)
{
    if (_type != AMF_STRICT_ARRAY) {
        return;
    }
    _value.array->push_back(val);
}

const AmfValue::MapType &AmfValue::GetMap() const
{
    return *_value.object;
}
const AmfValue::ArrayType &AmfValue::GetArr() const
{
    return *_value.array;
}

inline void AmfValue::Init()
{
    switch (_type) {
        case AMF_OBJECT:
        case AMF_ECMA_ARRAY:
            _value.object = new MapType;
            break;
        case AMF_STRING:
            _value.string = new std::string;
            break;
        case AMF_STRICT_ARRAY:
            _value.array = new ArrayType;
            break;

        default:
            break;
    }
}

inline void AmfValue::Destroy()
{
    switch (_type) {
        case AMF_STRING:
            if (_value.string) {
                delete _value.string;
                _value.string = nullptr;
            }
            break;
        case AMF_OBJECT:
        case AMF_ECMA_ARRAY:
            if (_value.object) {
                delete _value.object;
                _value.object = nullptr;
            }
            break;
        case AMF_STRICT_ARRAY:
            if (_value.array) {
                delete _value.array;
                _value.array = nullptr;
            }
            break;
        default:
            break;
    }
}

///////////////////////////////////////////////////////////////////////////

enum AMF0Type : uint8_t {
    AMF0_NUMBER = 0,
    AMF0_BOOLEAN,
    AMF0_STRING,
    AMF0_OBJECT,
    AMF0_MOVIECLIP,
    AMF0_NULL,
    AMF0_UNDEFINED,
    AMF0_REFERENCE,
    AMF0_ECMA_ARRAY,
    AMF0_OBJECT_END,
    AMF0_STRICT_ARRAY,
    AMF0_DATE,
    AMF0_LONG_STRING,
    AMF0_UNSUPPORTED,
    AMF0_RECORD_SET,
    AMF0_XML_OBJECT,
    AMF0_TYPED_OBJECT,
    AMF0_SWITCH_AMF3,
};

enum Amf3Type : uint8_t {
    AMF3_UNDEFINED,
    AMF3_NULL,
    AMF3_FALSE,
    AMF3_TRUE,
    AMF3_INTEGER,
    AMF3_NUMBER,
    AMF3_STRING,
    AMF3_LEGACY_XML,
    AMF3_DATE,
    AMF3_ARRAY,
    AMF3_OBJECT,
    AMF3_XML,
    AMF3_BYTE_ARRAY,
};

//-------------------- AMFEncoder --------------------//
AMFEncoder &AMFEncoder::operator<<(const char *s)
{
    // T: 1 B
    // L: 2 B
    // V: n B
    if (s) {
        buf += char(AMF0_STRING);
        auto len = strlen(s);
        assert(len <= 0xFFFF);
        uint16_t str_len = htons((uint16_t)len);
        buf.append((char *)&str_len, 2);
        buf += s;
    } else {
        buf += char(AMF0_NULL);
    }
    return *this;
}

AMFEncoder &AMFEncoder::operator<<(const std::string &s)
{
    if (!s.empty()) {
        buf += char(AMF0_STRING);
        assert(s.size() <= 0xFFFF);
        uint16_t str_len = htons((uint16_t)s.size());
        buf.append((char *)&str_len, 2);
        buf += s;
    } else {
        buf += char(AMF0_NULL);
    }
    return *this;
}

AMFEncoder &AMFEncoder::operator<<(std::nullptr_t)
{
    buf += char(AMF0_NULL);
    return *this;
}

AMFEncoder &AMFEncoder::WriteUndefined()
{
    buf += char(AMF0_UNDEFINED);
    return *this;
}

AMFEncoder &AMFEncoder::operator<<(const int n)
{
    return (*this) << (double)n;
}

AMFEncoder &AMFEncoder::operator<<(const double n)
{
    buf += char(AMF0_NUMBER);
    uint64_t encoded = SwapEndian64(*(uint64_t *)&n);
    buf.append((char *)&encoded, 8);
    return *this;
}

AMFEncoder &AMFEncoder::operator<<(const bool b)
{
    buf += char(AMF0_BOOLEAN);
    buf += char(b);
    return *this;
}

AMFEncoder &AMFEncoder::operator<<(const AmfValue &value)
{
    switch ((int)value.Type()) {
        case AMF_STRING:
            *this << value.AsString();
            break;
        case AMF_NUMBER:
            *this << value.AsNumber();
            break;
        case AMF_INTEGER:
            *this << value.AsInteger();
            break;
        case AMF_BOOLEAN:
            *this << value.AsBoolean();
            break;
        case AMF_OBJECT: {
            buf += char(AMF0_OBJECT);
            for (auto &pr : value.GetMap()) {
                WriteKey(pr.first);
                *this << pr.second;
            }
            WriteKey("");
            buf += char(AMF0_OBJECT_END);
        } break;
        case AMF_ECMA_ARRAY: {
            buf += char(AMF0_ECMA_ARRAY);
            uint32_t sz = htonl((uint32_t)value.GetMap().size());
            buf.append((char *)&sz, 4);
            for (auto &pr : value.GetMap()) {
                WriteKey(pr.first);
                *this << pr.second;
            }
            WriteKey("");
            buf += char(AMF0_OBJECT_END);
        } break;
        case AMF_NULL:
            *this << nullptr;
            break;
        case AMF_UNDEFINED:
            this->WriteUndefined();
            break;
        case AMF_STRICT_ARRAY: {
            buf += char(AMF0_STRICT_ARRAY);
            uint32_t sz = htonl((uint32_t)value.GetArr().size());
            buf.append((char *)&sz, 4);
            for (auto &val : value.GetArr()) {
                *this << val;
            }
        } break;
    }
    return *this;
}

void AMFEncoder::WriteKey(const std::string &s)
{
    assert(s.size() <= 0xFFFF);
    uint16_t str_len = htons((uint16_t)s.size());
    buf.append((char *)&str_len, 2);
    buf += s;
}

void AMFEncoder::Clear()
{
    buf.clear();
}

const std::string &AMFEncoder::Data() const
{
    return buf;
}

//-------------------- AMFDecoder --------------------//
uint8_t AMFDecoder::Front()
{
    if (pos_ >= (size_t)buffer_.Size()) {
        return 0;
    }
    return uint8_t(buffer_.Data()[pos_]);
}

uint8_t AMFDecoder::PopFront()
{
    if (version_ == 0 && Front() == AMF0_SWITCH_AMF3) {
        pos_++;
        version_ = 3;
    }

    if (pos_ >= (size_t)buffer_.Size()) {
        return 0;
    }
    return uint8_t(buffer_.Data()[pos_++]);
}

template <>
double AMFDecoder::Load<double>()
{
    if (PopFront() != AMF0_NUMBER) {
        return 0;
    }
    if (pos_ + 8 > (size_t)buffer_.Size()) {
        return 0;
    }

    uint64_t val = ((uint64_t)LoadBE32(&buffer_.Data()[pos_]) << 32) | LoadBE32(&buffer_.Data()[pos_ + 4]);
    double n = 0;
    static_assert(sizeof(n) == sizeof(val), "sizeof(double) not eq sizeof(uint64_t)");
    memcpy(&n, &val, sizeof(n));
    pos_ += 8;
    return n;
}

template <>
bool AMFDecoder::Load<bool>()
{
    if (PopFront() != AMF0_BOOLEAN) {
        return false;
    }
    return PopFront() != 0;
}
template <>
unsigned int AMFDecoder::Load<unsigned int>()
{
    unsigned int value = 0;
    for (int i = 0; i < 4; ++i) {
        uint8_t b = PopFront();
        if (i == 3) {
            /* use all bits from 4th byte */
            value = (value << 8) | b;
            break;
        }
        value = (value << 7) | (b & 0x7f);
        if ((b & 0x80) == 0) {
            break;
        }
    }
    return value;
}

template <>
int AMFDecoder::Load<int>()
{
    if (version_ == 3) {
        return (int)Load<unsigned int>();
    } else {
        return (int)Load<double>();
    }
}

template <>
std::string AMFDecoder::Load<std::string>()
{
    size_t str_len = 0;
    uint8_t type = PopFront();
    if (version_ == 3) {
        if (type != AMF3_STRING) {
            return {};
        }
        str_len = Load<unsigned int>() / 2;

    } else {
        if (type != AMF0_STRING) {
            return {};
        }
        if (pos_ + 2 > (size_t)buffer_.Size()) {
            return {};
        }
        str_len = LoadBE16(buffer_.Data() + pos_);
        pos_ += 2;
    }
    if (pos_ + str_len > (size_t)buffer_.Size()) {
        return {};
    }
    std::string s(buffer_.Peek() + pos_, str_len);
    pos_ += str_len;
    return s;
}

template <>
AmfValue AMFDecoder::Load<AmfValue>()
{
    uint8_t type = Front();
    if (version_ == 3) {
        switch (type) {
            case AMF3_STRING:
                return AmfValue(Load<std::string>());
            case AMF3_NUMBER:
                return AmfValue(Load<double>());
            case AMF3_INTEGER:
                return AmfValue(Load<int>());
            case AMF3_FALSE:
                pos_++;
                return AmfValue(false);
            case AMF3_TRUE:
                pos_++;
                return AmfValue(true);
            case AMF3_OBJECT:
                return LoadObject();
            case AMF3_ARRAY:
                return LoadEcma();
            case AMF3_NULL:
                pos_++;
                return AmfValue(AMF_NULL);
            case AMF3_UNDEFINED:
                pos_++;
                return AmfValue(AMF_UNDEFINED);
            default:
                return AmfValue();
        }
    } else {
        switch (type) {
            case AMF0_STRING:
                return AmfValue(Load<std::string>());
            case AMF0_NUMBER:
                return AmfValue(Load<double>());
            case AMF0_BOOLEAN:
                return AmfValue(Load<bool>());
            case AMF0_OBJECT:
                return LoadObject();
            case AMF0_ECMA_ARRAY:
                return LoadEcma();
            case AMF0_NULL:
                pos_++;
                return AmfValue(AMF_NULL);
            case AMF0_UNDEFINED:
                pos_++;
                return AmfValue(AMF_UNDEFINED);
            case AMF0_STRICT_ARRAY:
                return LoadArr();
            default:
                return AmfValue();
        }
    }
}

std::string AMFDecoder::LoadKey()
{
    if (pos_ + 2 > (size_t)buffer_.Size()) {
        return {};
    }
    size_t str_len = LoadBE16(buffer_.Data() + pos_);
    pos_ += 2;
    if (pos_ + str_len > (size_t)buffer_.Size()) {
        return {};
    }
    std::string s(buffer_.Peek() + pos_, str_len);
    pos_ += str_len;
    return s;
}

AmfValue AMFDecoder::LoadObject()
{
    AmfValue object(AMF_OBJECT);
    if (PopFront() != AMF0_OBJECT) {
        return AmfValue();
    }
    while (true) {
        std::string key = LoadKey();
        if (key.empty()) {
            break;
        }
        AmfValue value = Load<AmfValue>();
        object.Set(key, value);
    }
    if (PopFront() != AMF0_OBJECT_END) {
        return AmfValue();
    }
    return object;
}

AmfValue AMFDecoder::LoadEcma()
{
    /* ECMA array is the same as object, with 4 extra zero bytes */
    AmfValue object(AMF_ECMA_ARRAY);
    if (PopFront() != AMF0_ECMA_ARRAY) {
        return AmfValue();
    }
    if (pos_ + 4 > (size_t)buffer_.Size()) {
        return AmfValue();
    }
    pos_ += 4;
    while (true) {
        std::string key = LoadKey();
        if (key.empty()) {
            break;
        }
        AmfValue value = Load<AmfValue>();
        object.Set(key, value);
    }
    if (PopFront() != AMF0_OBJECT_END) {
        return AmfValue();
    }
    return object;
}

AmfValue AMFDecoder::LoadArr()
{
    /* ECMA array is the same as object, with 4 extra zero bytes */
    AmfValue object(AMF_STRICT_ARRAY);
    if (PopFront() != AMF0_STRICT_ARRAY) {
        return AmfValue();
    }
    if (pos_ + 4 > (size_t)buffer_.Size()) {
        return AmfValue();
    }

    uint32_t arrSize = SwapEndian32(*(uint32_t *)(buffer_.Data() + pos_));
    pos_ += 4;
    while (arrSize--) {
        AmfValue value = Load<AmfValue>();
        object.Add(value);
    }

    return object;
}

AMFDecoder::AMFDecoder(const DataBuffer &buf_in, size_t pos_in, int version_in) : version_(version_in), pos_(pos_in)
{
    buffer_.Assign(buf_in.Peek() + pos_in, (int)(buf_in.Size() - pos_in));
}
} // namespace Sharing
} // namespace OHOS